package org.example.Config;

import org.json.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.example.entity.ChatMessage;
import org.example.service.ChatMessageService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Collections;
import java.util.Date;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
public class ChatHandler extends TextWebSocketHandler {

    private static final Map<String, Set<WebSocketSession>> sessions = new ConcurrentHashMap<>();

    @Autowired
    private ChatMessageService chatHistoryRepository;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String project = getProjectFromSession(session);
        sessions.computeIfAbsent(project, k -> new CopyOnWriteArraySet<>()).add(session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String project = getProjectFromSession(session);
        String username = getUsernameFromSession(session);

        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setProjectName(project);
        chatMessage.setUsername(username);
        JSONObject jsonObject = new JSONObject(message.getPayload());

        // 提取 "message" 字段的值
        String messageContent = jsonObject.getString("message");

        chatMessage.setMessage(messageContent);
        chatMessage.setData(new Date());

        // 保存到数据库
        chatHistoryRepository.save(chatMessage);

        // 广播消息
        broadcast(project, chatMessage);
    }

    private void broadcast(String project, ChatMessage message) {
        ObjectMapper mapper = new ObjectMapper();
        sessions.getOrDefault(project, Collections.emptySet())
                .forEach(session -> {
                    try {
                        if (session.isOpen()) {
                            session.sendMessage(new TextMessage(mapper.writeValueAsString(message)));
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                });
    }

    private String getProjectFromSession(WebSocketSession session) {
        return session.getAttributes().get("project").toString();
    }

    private String getUsernameFromSession(WebSocketSession session) {
        return session.getAttributes().get("username").toString();
    }
}